package com.agentframework.examples.playwright;

import com.agentframework.core.tools.AgentTool;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.microsoft.playwright.*;
import com.microsoft.playwright.options.WaitForSelectorState;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;

import java.nio.file.Path;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * @author wangjun
 **/
public class PlaywrightTool implements AgentTool {

    protected static final Logger logger = LoggerFactory.getLogger(PlaywrightTool.class);

    private final Playwright playwright;

    private Browser browser;

    private BrowserContext context;

    private Page page;

    private final boolean headless;

    final boolean returnDirect = true;

    public PlaywrightTool(boolean headless) {
        this.playwright = Playwright.create();
        this.headless = headless;
    }

    public PlaywrightTool() {
        this(false);
    }

    @Tool(description = "Start playwright", returnDirect = returnDirect)
    public void start() {
        this.browser = playwright.chromium().launch(new BrowserType.LaunchOptions().setHeadless(headless));
        createNewContext();
    }

    @Tool(description = "Create new browser context with isolation", returnDirect = returnDirect)
    public void createNewContext() {
        if (context != null) {
            context.close();
        }
        this.context = browser.newContext();
        this.page = context.newPage();
    }

    @Tool(description = "Navigate to specified URL", returnDirect = returnDirect)
    public void navigateTo(@ToolParam(description = "Full URL starting with https:// or http://") String url) {
        page.navigate(url);
    }

    @Tool(description = "Go back in browser history", returnDirect = returnDirect)
    public void goBack() {
        page.goBack();
    }

    @Tool(description = "Go forward in browser history", returnDirect = returnDirect)
    public void goForward() {
        page.goForward();
    }

    @Tool(description = "Upload file to file input element", returnDirect = returnDirect)
    public void chooseFile(
            @ToolParam(description = "iframe url or id or name") String frame,
            @ToolParam(description = "Xpath or CSS selector for file input") String selector,
            @ToolParam(description = "Absolute path to the file") Path filePath) {
        getLocator(selector, frame).setInputFiles(filePath);
    }

    @Tool(description = "Capture page screenshot", returnDirect = returnDirect)
    public String takeScreenshot(
            @ToolParam(description = "Output file path") Path savePath,
            @ToolParam(description = "Capture full scrollable page") boolean fullPage) {
        page.screenshot(new Page.ScreenshotOptions()
                .setPath(savePath)
                .setFullPage(fullPage));
        return savePath.toString();
    }

    @Tool(description = "Click page element", returnDirect = returnDirect)
    public void click(
            @ToolParam(description = "iframe url or id or name") String frame,
            @ToolParam(description = "Xpath or CSS selector of target element") String selector) {
        getLocator(selector, frame).click();
    }

    @Tool(description = "Hover over element", returnDirect = returnDirect)
    public void hover(@ToolParam(description = "Xpath or CSS selector of target element") String selector) {
        page.locator(selector).hover();
    }

    @Tool(description = "Type text into input field", returnDirect = returnDirect)
    public void type(
            @ToolParam(description = "iframe url or id or name") String frame,
            @ToolParam(description = "Xpath or CSS selector of input element") String selector,
            @ToolParam(description = "Text content to input") String text) {
        getLocator(selector, frame).fill(text);
    }

    @Tool(description = "Select dropdown option", returnDirect = returnDirect)
    public void selectOption(
            @ToolParam(description = "iframe url or id or name") String frame,
            @ToolParam(description = "Xpath or CSS selector of dropdown") String selector,
            @ToolParam(description = "Value of option to select") String value) {
        getLocator(selector, frame).selectOption(value);
    }

    @Tool(description = "Simulate keyboard press", returnDirect = returnDirect)
    public void pressKey(@ToolParam(description = "Key name (e.g. Enter, Tab)") String key) {
        page.keyboard().press(key);
    }

    @Tool(description = "Wait for specified duration", returnDirect = returnDirect)
    public void waitFor(@ToolParam(description = "Wait time in milliseconds") long milliseconds) {
        page.waitForTimeout(milliseconds);
    }

    @Tool(description = "Save page as PDF document", returnDirect = returnDirect)
    public void saveAsPdf(@ToolParam(description = "Output file path") Path savePath) {
        page.pdf(new Page.PdfOptions().setPath(savePath));
    }

    @Tool(description = "Close browser and cleanup resources", returnDirect = returnDirect)
    public void close() {
        if (context != null) context.close();
        if (browser != null) browser.close();
        playwright.close();
    }

    // Enhanced methods
    @Tool(description = "Wait for element visibility", returnDirect = returnDirect)
    public void waitForElementVisible(String selector, int timeoutSeconds) {
        long timeout = TimeUnit.SECONDS.toMillis(timeoutSeconds);
        page.waitForSelector(selector,
                new Page.WaitForSelectorOptions()
                        .setState(WaitForSelectorState.VISIBLE)
                        .setTimeout(timeout));
    }

    @Tool(description = "Enable network traffic monitoring", returnDirect = returnDirect)
    public void enableNetworkMonitoring() {
        context.onRequest(request ->
                logger.debug("Request: method: {}, url: {} ", request.method(), request.url()));
        context.onResponse(response ->
                logger.debug("Response: status: {}, url: {} ", response.status(), response.url()));
    }

    @Tool(description = "Get current page URL", returnDirect = returnDirect)
    public String getCurrentUrl() {
        return page.url();
    }

    @Tool(description = "Get page title", returnDirect = returnDirect)
    public String getPageTitle() {
        return page.title();
    }

    @Tool(description = "Get page content", returnDirect = returnDirect)
    public String getPageContent() throws Exception {
        List<Map<String, Object>> elementList = new ArrayList<>();
        extractDomElements(page.mainFrame(), page.url(), elementList);
        String result = new ObjectMapper().writeValueAsString(elementList);
        logger.info("page content: {}", result);
        return result;
    }

    private Locator getLocator(String selector, String iframe) {
        if (iframe == null || iframe.isEmpty()) {
            return page.locator(selector);
        }
        for (Frame frame : page.frames()) {
            String id = frame.name();
            String url = frame.url();
            if (iframe.equalsIgnoreCase(id) || url.contains(iframe)) {
                return frame.locator(selector);
            }
        }
        return page.locator(selector);
    }

    private void extractDomElements(Frame frame, String frameUrl, List<Map<String, Object>> collector) {
        // 1. 交互式选择器
        String selector = "input, button, textarea, select, a, div, ul, li";

        List<ElementHandle> elements = frame.querySelectorAll(selector);

        for (ElementHandle element : elements) {
            try {
                Map<String, Object> info = new HashMap<>();

                // 标签名
                String tag = (String) element.evaluate("el => el.tagName.toLowerCase()");
                info.put("tag", tag);

                // XPath
                info.put("xpath", generateSmartXPath(element));

                // frame URL
                info.put("frameUrl", frameUrl);

                // type 属性（如 input 类型）
                String type = element.getAttribute("type");
                if (type != null) info.put("type", type);

                // placeholder 属性
                String placeholder = element.getAttribute("placeholder");
                if (placeholder != null) info.put("placeholder", placeholder);

                // 按钮类元素文本内容
                if ("button".equals(tag)) {
                    String text = element.innerText();
                    info.put("text", text);
                }

                collector.add(info);
            } catch (Exception e) {
                logger.error("解析元素失败：{}", e.getMessage());
            }
        }

        // 2. 遍历 iframe（递归）
        List<ElementHandle> iframeHandles = frame.querySelectorAll("iframe");

        for (ElementHandle iframeHandle : iframeHandles) {
            try {
                Frame childFrame = iframeHandle.contentFrame();
                if (childFrame != null) {
                    extractDomElements(childFrame, childFrame.url(), collector);
                }
            } catch (Exception e) {
                logger.error("无法获取 iframe 内容，可能是跨域限制：{}", e.getMessage());
            }
        }
    }

    private String generateSmartXPath(ElementHandle element) {
        try {
            String id = element.getAttribute("id");
            if (id != null && !id.isEmpty()) {
                return "//*[@id=\"" + id + "\"]";
            }

            String name = element.getAttribute("name");
            if (name != null && !name.isEmpty()) {
                return "//*[@name=\"" + name + "\"]";
            }

            String type = element.getAttribute("type");
            String placeholder = element.getAttribute("placeholder");
            if (type != null && placeholder != null) {
                return "//input[@type=\"" + type + "\" and @placeholder=\"" + placeholder + "\"]";
            }

            String tag = (String) element.evaluate("el => el.tagName.toLowerCase()");
            String textContent = (String) element.evaluate("el => el.textContent?.trim()");
            if (tag.equals("button") && textContent != null && !textContent.isEmpty()) {
                return "//button[contains(text(),\"" + textContent + "\")]";
            }

            // 获取路径格式的xpath
            return (String) element.evaluate("el => {" +
                    "function getXPath(el) {" +
                    "  const parts = [];" +
                    "  while (el && el.nodeType === 1) {" +
                    "    let index = 1;" +
                    "    let sibling = el.previousSibling;" +
                    "    while (sibling) {" +
                    "      if (sibling.nodeType === 1 && sibling.tagName === el.tagName) index++;" +
                    "      sibling = sibling.previousSibling;" +
                    "    }" +
                    "    parts.unshift(el.tagName.toLowerCase() + '[' + index + ']');" +
                    "    el = el.parentNode;" +
                    "  }" +
                    "  return '/' + parts.join('/');" +
                    "}" +
                    "return getXPath(el);" +
                    "}");
        } catch (Exception e) {
            logger.error("XPath 生成失败：{}", e.getMessage());
            return null;
        }
    }

}
